#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import QuaternionStamped, PointStamped, Vector3Stamped
from nav_msgs.msg import Odometry
from tf_transformations import quaternion_from_euler
import math
import yaml
import os
from pathlib import Path

class WaypointNavigation(Node):
    def __init__(self):
        super().__init__('waypoint_navigation')
        
        # 创建发布器
        self.attitude_pub_ = self.create_publisher(QuaternionStamped, '/bluerov00/attitude_target_desired', 10)
        self.position_pub_ = self.create_publisher(PointStamped, '/bluerov00/position_target', 10)
        
        # 订阅当前位姿
        self.odometry_sub_ = self.create_subscription(
            Odometry, 
            '/bluerov00/odometry', 
            self.odometry_callback, 
            10
        )
        
        # 加载配置文件
        config_path = '/home/siyuan/ros2/src/hippo_control/scripts/bernoulli_default.yaml'  # 假设配置文件在同一目录
        self.waypoints = self.load_waypoints(config_path)
        
        # 导航状态
        self.current_waypoint_index = 0
        self.current_pose = None
        self.reached_current_waypoint = False
        
        self.threshold = 0.6

        # 定时器
        timer_period = 0.1  # 10Hz
        self.timer = self.create_timer(timer_period, self.navigation_loop)
        
        self.get_logger().info(f'Loaded {len(self.waypoints)} waypoints')
        self.publish_current_waypoint()
    
    def load_waypoints(self, config_path):
        """从YAML文件加载路径点配置"""
        try:
            with open(config_path, 'r') as file:
                config = yaml.safe_load(file)
            
            waypoints = config.get('waypoints', [])
            self.loop = config.get('loop', False)
            
            self.get_logger().info(f'Successfully loaded {len(waypoints)} waypoints from {config_path}')
            return waypoints
            
        except Exception as e:
            self.get_logger().error(f'Failed to load waypoints from {config_path}: {e}')
            return []
    
    def odometry_callback(self, msg):
        """更新当前位姿"""
        self.current_pose = msg.pose.pose
    
    def calculate_distance(self, pose1, pose2):
        """计算两个位姿之间的欧几里得距离"""
        dx = pose1.position.x - pose2.position.x
        dy = pose1.position.y - pose2.position.y
        dz = pose1.position.z - pose2.position.z
        return math.sqrt(dx**2 + dy**2 + dz**2)
    
    def publish_attitude_target(self):
        """发布姿态目标"""
        msg = QuaternionStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'map'
        
        waypoint = self.waypoints[self.current_waypoint_index]
        yaw_deg = waypoint['yaw']
        pitch_deg = waypoint['pitch']
        roll_deg = waypoint['roll']
        # 将角度转换为弧度
        yaw_rad = yaw_deg * math.pi / 180.0
        pitch_rad = pitch_deg * math.pi / 180.0
        roll_rad = roll_deg * math.pi / 180.0
        
        # 只控制偏航角，横滚和俯仰保持为0
        quaternion = quaternion_from_euler(roll_rad, pitch_rad, yaw_rad)
        
        msg.quaternion.x = quaternion[0]
        msg.quaternion.y = quaternion[1]
        msg.quaternion.z = quaternion[2]
        msg.quaternion.w = quaternion[3]
        
        self.attitude_pub_.publish(msg)

        self.get_logger().info(
            f'Waypoint {self.current_waypoint_index + 1}/{len(self.waypoints)}: '
            f'Yaw: {waypoint["yaw"]:.2f}, Pitch: {waypoint["pitch"]:.2f}, Roll: {waypoint["roll"]:.2f}'
        )
    
    def publish_current_waypoint(self):
        """发布当前目标路径点"""
        if self.current_waypoint_index >= len(self.waypoints):
            return
            
        waypoint = self.waypoints[self.current_waypoint_index]
        
        msg = PointStamped()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'world'
        
        # 关键修复：确保使用float类型
        msg.point.x = float(waypoint['x'])
        msg.point.y = float(waypoint['y'])
        msg.point.z = float(waypoint['z'])
        
        self.position_pub_.publish(msg)
        
        self.get_logger().info(
            f'Waypoint {self.current_waypoint_index + 1}/{len(self.waypoints)}: '
            f'X: {waypoint["x"]:.2f}m, Y: {waypoint["y"]:.2f}m, Z: {waypoint["z"]:.2f}m'
        )
    
    def check_waypoint_reached(self):
        """检查是否到达当前路径点"""
        if self.current_pose is None or self.current_waypoint_index >= len(self.waypoints):
            return False
        
        waypoint = self.waypoints[self.current_waypoint_index]
        
        # 创建目标位姿
        target_pose = type('', (), {})()
        target_pose.position = type('', (), {})()
        target_pose.position.x = waypoint['x']
        target_pose.position.y = waypoint['y'] 
        target_pose.position.z = waypoint['z']
        
        # 计算距离
        distance = self.calculate_distance(self.current_pose, target_pose)
        # threshold = 1
        
        return distance <= self.threshold
    
    def move_to_next_waypoint(self):
        """移动到下一个路径点"""
        self.current_waypoint_index += 1
        
        if self.current_waypoint_index >= len(self.waypoints):
            if self.loop:
                self.current_waypoint_index = 0
                self.get_logger().info('Looping back to first waypoint')
            else:
                self.get_logger().info('All waypoints completed!')
                return False
        
        self.publish_current_waypoint()
        return True
    
    def navigation_loop(self):
        """主导航循环"""
        if not self.waypoints:
            self.get_logger().warning('No waypoints loaded')
            return
        
        if self.current_pose is None:
            self.get_logger().warning('Waiting for odometry data...')
            return
        
        if self.current_waypoint_index >= len(self.waypoints):
            return
        
        # 检查是否到达当前路径点
        if self.check_waypoint_reached():
            if not self.reached_current_waypoint:
                self.get_logger().info(f'Reached waypoint {self.current_waypoint_index + 1}')
                self.reached_current_waypoint = True
                
                # 短暂停留后移动到下一个点
                if not self.move_to_next_waypoint():
                    return
        else:
            self.reached_current_waypoint = False
        

        self.publish_attitude_target()
        
        # 发布当前位置目标
        self.publish_current_waypoint()
        
        # 显示导航状态
        waypoint = self.waypoints[self.current_waypoint_index]
        distance = self.calculate_distance(
            self.current_pose, 
            type('', (), {
                'position': type('', (), {
                    'x': waypoint['x'],
                    'y': waypoint['y'], 
                    'z': waypoint['z']
                })()
            })()
        )
        
        self.get_logger().info(
            f'Navigating to waypoint {self.current_waypoint_index + 1}, '
            f'Distance: {distance:.2f}m', 
            throttle_duration_sec=1.0  # 每秒只显示一次
        )

def main():
    rclpy.init()
    node = WaypointNavigation()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()